package org.example.brave.dubbo.consumer.filter;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import zipkin2.reporter.AsyncReporter;
import zipkin2.reporter.brave.ZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

import java.util.Map;

@Activate
public class BraveConsumerFilter implements Filter {
    public static OkHttpSender sender;
    public static Tracing tracing;

    static {
        sender = OkHttpSender.newBuilder().endpoint("http://localhost:9411/api/v2/spans").build();
        tracing = Tracing.newBuilder().localServiceName("brave-dubbo-consumer")
                .addSpanHandler(ZipkinSpanHandler.create(AsyncReporter.builder(sender).build()))
                .build();
    }

    Propagation.RemoteSetter<Map<String, Object>> SETTER = new Propagation.RemoteSetter<Map<String, Object>>() {
        @Override
        public Span.Kind spanKind() {
            return Span.Kind.CLIENT;
        }

        @Override
        public void put(Map<String, Object> clientAttachment, String key, String value) {
            clientAttachment.put(key, value);
        }
    };

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        System.out.println("Consumer judge whether provider or consumer: " + (isConsumer(invocation) ? "consumer" : "provider"));
        return process(invoker, invocation);
    }

    private boolean isConsumer(Invocation invocation) {
        Invoker<?> invoker = invocation.getInvoker();
        return invoker.getUrl().getParameter("side", "provider").equals("consumer");
    }

    private Result process(Invoker<?> invoker, Invocation invocation) {
        Result result = null;
        TraceContext.Injector<Map<String, Object>> injector = tracing.propagation().injector(SETTER);
        CurrentTraceContext currentTraceContext = tracing.currentTraceContext();
        Tracer tracer = tracing.tracer();
        Span span = tracer.nextSpan().name(invocation.getMethodName()).start();
        try (CurrentTraceContext.Scope ws = currentTraceContext.newScope(span.context())) {
            injector.inject(span.context(), RpcContext.getClientAttachment().getObjectAttachments());
            span.kind(Span.Kind.CLIENT);
            System.out.println("consumer filter test...");
            result = invoker.invoke(invocation);
        } catch (Exception e) {
            span.error(e);
            e.printStackTrace();
        } finally {
            span.finish();
        }
        return result;
    }
}
